from utils.metrics import batch_bleu, batch_acc, batch_rouge, batch_meteor
from argparse import ArgumentParser
import os
from utils.jsonl import JsonLReader
import csv
from tqdm import tqdm


def bleu1(batch_hyp, batch_ref):
    return batch_bleu(batch_hyp, batch_ref, n=1)


def bleu2(batch_hyp, batch_ref):
    return batch_bleu(batch_hyp, batch_ref, n=2)


def bleu3(batch_hyp, batch_ref):
    return batch_bleu(batch_hyp, batch_ref, n=3)


def bleu4(batch_hyp, batch_ref):
    return batch_bleu(batch_hyp, batch_ref, n=4)


def bleu4_smoothed(batch_hyp, batch_ref):
    return batch_bleu(batch_hyp, batch_ref, n=4, smooth_method=3)


def meteor(batch_hyp, batch_ref):
    return batch_meteor(batch_hyp, batch_ref)


def rouge(batch_hyp, batch_ref):
    return batch_rouge(batch_hyp, batch_ref)


def acc(batch_hyp, batch_ref):
    return batch_acc(batch_hyp, batch_ref)


def eval_file(file_path):
    data = JsonLReader(file_path).read()
    hyp = [each['hyp'] for each in data]
    ref = [each['ref'] for each in data]
    return [
        file_path.split("/")[-1],
        bleu1(hyp, ref),
        bleu2(hyp, ref),
        bleu3(hyp, ref),
        bleu4(hyp, ref),
        bleu4_smoothed(hyp, ref),
        meteor(hyp, ref),
        rouge(hyp, ref),
        acc(hyp, ref),
    ]


def main(args):
    dir_path = args.dir
    files = os.listdir(dir_path)
    heads = [
        'file',
        'bleu1',
        'bleu2',
        'bleu3',
        'bleu4',
        'bleu4-smoothed',
        'meteor',
        'rouge-l',
        'acc',
    ]
    result = [

    ]

    for each in tqdm(files):
        file_path = os.path.join(dir_path, each)
        result.append(eval_file(file_path))
    f = open(args.out, "w+")
    f = csv.writer(f)
    f.writerow(heads)
    f.writerows(result)


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument("--dir", help="", type=str, default="./translate")
    parser.add_argument("--out", help="", type=str, default="./translate/results.csv")
    cli_args = parser.parse_args()
    main(cli_args)
